import datetime
from matplotlib import pyplot as plt
import pandas as pd
from PIL import Image
from utils.consts import LOG_PATH, DOC_PATH

PIC_PATH = 'imgs'
X_Y_FONT = 12
LEGEND_FONT = 14
TITLE_FONT = 22
FIGURE_SIZE = (21, 9)
DPI = 300
X_INDEX_NUM = 30
Y_INDEX_NUM = 10
X_ROTATE = 45
COLORS = ['#000075', '#f58231', '#e6194B', '#3cb44b', '#ffe119', '#911eb4', '#42d4f4', '#f032e6', '#bfef45',
          '#fabed4', '#469990', '#dcbeff', '#9A6324', '#800000', '#808000']
EDGE_RATE = 0.03
IMAGE_SIZE = 1000


class PointSet:
    def __init__(self, list_x, list_y, name, win_size=19) -> None:
        self.x = list_x
        self.y = PointSet.moving_average(list_y, win_size)
        self.name = name

    def moving_average(data: list[float], win_size) -> list[float]:
        res_data = [0.]*len(data)
        for i in range(len(data)):
            left = i - win_size if i >= win_size else 0
            right = i if i != 0 else 1
            res_data[i] = sum(data[left:i])/(right-left)
        return res_data


def draw_plot(datas: list[PointSet], title: str, path: str, best_step: int, label_pos: int = 1):
    x_limit = get_bound([ps.x for ps in datas])
    y_limit = get_bound([ps.y for ps in datas])
    x_range = x_limit[1] - x_limit[0]
    y_range = y_limit[1] - y_limit[0]
    x_tick_list = [x_limit[0] + i * x_range /
                   X_INDEX_NUM for i in range(X_INDEX_NUM)]
    y_tick_list = [y_limit[0] + i * y_range /
                   Y_INDEX_NUM for i in range(Y_INDEX_NUM)]
    plt.subplots(1, 1, figsize=FIGURE_SIZE, dpi=DPI)

    # draw plot
    for i, ps in enumerate(datas):
        plt.plot(ps.x, ps.y, lw=1.5, color=COLORS[i], label=ps.name)

    # draw y tick line
    plt.hlines(y_tick_list, x_limit[0], x_limit[1], colors='black',
               alpha=0.3, linestyles='--', lw=0.5)

    # draw model saved step line
    plt.vlines(best_step, y_limit[0], y_limit[1],
               colors='black', alpha=0.3, linestyles='dashdot', lw=0.5)

    # decorate
    plt.tick_params(axis='both', which='both', bottom=False, top=False,
                    labelbottom=True, left=False, right=False, labelleft=True)

    # Lighten borders
    plt.gca().spines["top"].set_alpha(.3)
    plt.gca().spines["bottom"].set_alpha(.3)
    plt.gca().spines["right"].set_alpha(.3)
    plt.gca().spines["left"].set_alpha(.3)

    plt.title(title, fontsize=TITLE_FONT)
    plt.yticks(y_tick_list, [
               f'{y:.4f}' for y in y_tick_list], fontsize=X_Y_FONT)
    plt.xticks(x_tick_list, [f'{float(x):.0f}' for x in x_tick_list],
               horizontalalignment='right', fontsize=X_Y_FONT, rotation=X_ROTATE)
    plt.ylim(y_limit[0]-y_range*EDGE_RATE, y_limit[1]+y_range*EDGE_RATE)
    plt.xlim(x_limit[0]-x_range*EDGE_RATE, x_limit[1]+x_range*EDGE_RATE)
    plt.legend(loc=label_pos, fontsize=LEGEND_FONT,
               ncol=round(len(datas)**0.5))

    plt.savefig(path)
    plt.close('all')


def get_bound(inputs: list[list]):
    maxs = [max(l) for l in inputs]
    mins = [min(l) for l in inputs]
    return min(mins), max(maxs)


def draw_bar(distr: list, cls_dict: list[str], title: str, path: str):
    _, ax = plt.subplots(1, 1, figsize=FIGURE_SIZE, dpi=DPI)

    num_cls = len(cls_dict)

    if isinstance(distr[0], int):
        distr = [distr]

    dataX = [i for i in range(num_cls)]
    intv = 0.35
    y_max = distr[0][0]
    for i, dataY in enumerate(distr):
        x_pos = [dataX[j] + intv*i for j in range(num_cls)]
        plt.vlines(x=x_pos, ymin=0, ymax=dataY,
                   color=COLORS[i], alpha=0.7, linewidth=40)
        # Annotate Text
        for i in range(num_cls):
            ax.text(x_pos[i], dataY[i]+y_max*0.01, dataY[i],
                    horizontalalignment='center')

    # decorate
    plt.tick_params(axis='both', which='both', bottom=False, top=False,
                    labelbottom=True, left=False, right=False, labelleft=True)

    # Lighten borders
    plt.gca().spines["top"].set_alpha(.3)
    plt.gca().spines["bottom"].set_alpha(.3)
    plt.gca().spines["right"].set_alpha(.3)
    plt.gca().spines["left"].set_alpha(.3)

    plt.title(title, fontsize=TITLE_FONT)
    dataX = [dataX[i] + intv*(len(distr)-1)/2 for i in range(num_cls)]
    plt.xticks(dataX, [f'{l.upper()}' for l in cls_dict],
               horizontalalignment='center', fontsize=X_Y_FONT)
    plt.yticks(range(0, y_max+1, round(y_max/10)),
               range(0, y_max+1, round(y_max/10)), fontsize=X_Y_FONT)
    plt.ylim(0, y_max*1.1)

    plt.savefig(path)
    plt.close('all')


def draw_image(path: str, cls_dict: list[str]):
    gen = Image.open(path).resize((IMAGE_SIZE, IMAGE_SIZE))
    itv = IMAGE_SIZE / len(cls_dict)
    pic_id = range(len(cls_dict))
    pos = [(i+0.5)*itv for i in pic_id]
    plt.figure(1, dpi=DPI)
    plt.imshow(gen)
    plt.title('Generate Images', fontsize=14)
    plt.xticks(pos, pic_id, fontsize=10)
    plt.yticks(pos, cls_dict, fontsize=10)
    plt.savefig(path)
    plt.close('all')


def draw_matrix(path: str, rp_mat: list[list[int]], cls_dict: list[str]):
    fig = plt.figure(dpi=DPI)
    axs = fig.add_subplot()
    caxes = axs.matshow(rp_mat, interpolation='nearest')
    cb = fig.colorbar(caxes)
    cb.ax.tick_params(labelsize=8)
    plt.title('Confusion Matrix', fontsize=12)
    axs.set_xlabel('Real Label', fontsize=10)
    axs.set_ylabel('Prediction', fontsize=10)
    axs.set_xticks(range(0, len(cls_dict)), cls_dict, fontsize=8, rotation=30)
    axs.set_yticks(range(0, len(cls_dict)), cls_dict, fontsize=8)
    plt.subplots_adjust(top=0.85)
    plt.savefig(path)
    plt.close('all')


def draw(taskid: str, conf: dict):
    task = taskid[0]
    df = pd.read_csv(f"{LOG_PATH}/{taskid}.csv")
    d_x = df['total_step']
    save_path = f'{DOC_PATH}/{PIC_PATH}/{taskid}'
    best_step = conf['best_step']
    cls_dict = conf['cls_dict']
    distr = conf['distr']
    if task == 'A' or task == 'B':
        ps_loss = [PointSet(d_x, df[k].values, k, 5)
                   for k in ['tr_loss', 'vl_loss']]
        draw_plot(ps_loss, 'Loss-Figure', f'{save_path}-loss.png', best_step)

        ps_metrics = [PointSet(d_x, df[k].values, k, 5)
                      for k in ['acc', 'mbr', 'mbp']]
        draw_plot(ps_metrics, 'Metrics-Figure',
                  f'{save_path}-metrics.png', best_step, 4)

        recall_keys = [k for k in df.columns if k.startswith('recall_')]
        ps_accs = [PointSet(d_x, df[k].values, k, 19) for k in recall_keys]
        draw_plot(ps_accs, 'Recall-Figure',
                  f'{save_path}-recall.png', best_step, 4)

        precision_keys = [k for k in df.columns if k.startswith('precision_')]
        ps_accs = [PointSet(d_x, df[k].values, k, 19) for k in precision_keys]
        draw_plot(ps_accs, 'Precision-Figure',
                  f'{save_path}-precision.png', best_step, 4)

        draw_matrix(f'{save_path}-matrix.png', conf['rp_mat'], cls_dict)
    else:
        ps_loss = [PointSet(d_x, df[k].values, k, 19)
                   for k in ['LossD', 'LossG']]
        draw_plot(ps_loss, 'Loss-Figure', f'{save_path}-loss.png', best_step)

        ps_loss = [PointSet(d_x, df[k].values, k, 19)
                   for k in ['D(x)', 'D(G0(z))', 'D(G1(z))']]
        draw_plot(ps_loss, 'D(x), D(G0(z)), D(G1(z)) Figure',
                  f'{save_path}-disc.png', best_step)

        draw_image(f'{save_path}-gen.png', cls_dict)

    draw_bar(distr, cls_dict, 'Trainset Distribution Figure',
             f'{save_path}-distr.png')


def gen_doc(taskid: str, conf: dict):
    task_id = taskid
    timestamp = task_id.split('-')[-1]
    time = datetime.datetime.fromtimestamp(
        int(timestamp)).strftime("%Y-%m-%d %H:%M:%S")

    task_msg = '...'
    figure_msg = '...'

    def Line(text: str, prefix: str = '-') -> str:
        return f'{prefix} {text}\n\n'

    def Title(text: str, level: int = 2) -> str:
        hashtags = '#'*level
        return f'{hashtags} {text}\n\n'

    def Separator():
        return '\n----\n'

    def Img(name: str) -> str:
        return f'![img]({PIC_PATH}/{task_id}-{name})\n\n'

    def Sequence(lines: list[str]):
        return '\n'.join(lines)

    if conf['task'] in ['stageA', 'stageB']:
        task_msg = Sequence([
            Title(f"{conf['task'].upper()} Settings:"),
            Line(f"Backbone: {conf['backbone']}")
        ])
        figure_msg = Sequence([
            Title("Training Process:"),
            Img("loss.png"),
            Separator(),
            Title("Experiment Outcomes:"),
            Line(f"Best Accuracy: {conf['max_acc']:.6f}"),
            Line(f"Max Bottom Precision: {conf['max_bottom_precision']:.6f}"),
            Line(f"Max Bottom Recall: {conf['max_bottom_recall']:.6f}"),
            Line("Figures:"),
            Img("metrics.png"),
            Img("recall.png"),
            Img("precision.png"),
            Img("matrix.png")
        ])
    if conf['task'] == 'stageB':
        gan_path = conf['gan'] if 'gan' in conf.keys() else '(NONE)'
        task_msg = Sequence([
            task_msg,
            Line(f"Backbone Pre-Trianed Model: {conf['model']}"),
            Line(f"Augment Methods: {conf['method']}"),
            Line(f"GAN Model: {gan_path}")
        ])
    if conf['task'] == 'stageG':
        task_msg = Sequence([
            Title(f"{conf['task'].upper()} Settings:"),
            Line(
                f"Net Work: cDC-GAN(Conditional Deep Convolution Generative Adversarial Networks)")
        ])
        figure_msg = Sequence([
            Title("Training Process:"),
            Img("loss.png"),
            Img("disc.png"),
            Separator(),
            Title("Generate Pics:"),
            Img("gen.png")
        ])

    msg = Sequence([
        Title(f"Exp: {task_id}", 1),
        Line(f"Time: {time}", '>'),
        Line("", '>'),
        Line(f"Task Type: {conf['task']}", '>'),
        Separator(),
        Title("Basic Message:"),
        Line(f"Dataset: {conf['dataset'].upper()}"),
        Line(f"Numcls: {len(conf['cls_dict'])}"),
        Line(f"Distribution: {conf['distr']}"),
        Line(f"Classes Dict: {conf['cls_dict']}"),
        Img("distr.png"),
        Separator(),
        Title("Training Configs:"),
        Line(f"Total Epochs: {conf['epoch']}"),
        Line(f"Batch Size: {conf['batch']}"),
        Line(f"Learning Rate: {conf['lr']:.6f}"),
        Separator(),
        task_msg,
        Separator(),
        figure_msg
    ])

    with open(f'{DOC_PATH}/{task_id}.md', 'w') as f:
        f.write(msg)


if __name__ == '__main__':
    import argparse
    import warnings
    from utils.tools import read_json, check_path
    warnings.filterwarnings("ignore", category=DeprecationWarning)

    parser = argparse.ArgumentParser(description='data path')
    parser.add_argument('--taskid', '-t', type=str, required=True)
    args = parser.parse_args()
    taskid = args.taskid

    conf_dict = read_json(f'{LOG_PATH}/{taskid}.json')
    check_path(DOC_PATH)
    check_path(f'{DOC_PATH}/{PIC_PATH}')

    draw(taskid, conf_dict)
    gen_doc(taskid, conf_dict)
